#include <test/testMyData.cuh>
#include <utils/DataConvert.cuh>
#include <iostream>

#include "DataCatcher.h"
#include "Profiler.h"

int testMyData(int argc, char* argv[])
{
    std::string inputFile("Vely-1.csr");

    std::cout << "Loading input file: " << inputFile << std::endl;

    MPI_Init(&argc, &argv);

    // AMG Param
    int max_levels = 20;
    int coarse_size = 8;
    double strength_threshold = 0.25;
    int aggressive_levels = 50;
    int coarsen_type = 0; // 0: HMIS, 1: PMIS
    int interp_type = 0; // 0: Long-Range Interpolation, 1: Smoothed Aggregation
    int interp_min_elements = 4;
    int interp_max_elements = 4;
    double truncation_factor = 0.1;
    double sparsification_threshold = 0.01;
    int cycle_type = 0; // 0: V-Cycle, 1: W-Cycle, 2: F-Cycle
    int pre_sweeps = 1;
    int post_sweeps = 1;
    int coarse_sweeps = 2;
    int smooth_type = 0; // 0: Jacobi, 1: SOR, 2: ILU, 3: Chebyshev

    double jacobi_factor = 0.75;

    int sor_type = 2; // 0: Forward, 1: Backword, 2: Symmetric
    double sor_factor = 1.0;

    int ilu_maxfil = 0;
    double ilu_droptol = 0.01;

    int chebyshev_order = 4;
    double chebyshev_eigen_ratio = 0.3;

    int gmres_restart = 10;
    int max_iters = 500;
    double tolerance = 1.0e-08;

    bool print_stats = 1;

    MPI_Comm comm = MPI_COMM_WORLD;

    auto readA = myamg::CSR2YHAMGCsrMatrix(inputFile);
    auto readB = myamg::CSRRhs2YHAMGVector(inputFile);

    auto A = std::make_unique<YHAMG::ParCSRMatrix>(*readA);
    auto b = std::make_unique<YHAMG::ParVector>(*readB);
    auto x = std::make_unique<YHAMG::ParVector>(comm);
    x->Resize(b->local.size);
    x->Fill(0.0);

    auto P = std::make_unique<YHAMG::ParCSRPrecondAMG>();

    P->MaxLevels = max_levels;
    P->CoarseSize = coarse_size;
    P->StrengthThreshold = strength_threshold;
    P->AggressiveLevels = aggressive_levels;
    P->CoarsenType = coarsen_type;
    P->InterpType = interp_type;
    P->InterpMinElements = interp_min_elements;
    P->InterpMaxElements = interp_max_elements;
    P->TruncationFactor = truncation_factor;
    P->SparsificationThreshold = sparsification_threshold;
    P->CycleType = cycle_type;
    P->PreSweeps = pre_sweeps; // pre步数
    P->PostSweeps = post_sweeps; // post步数
    P->CoarseSweeps = coarse_sweeps; // 最后一层步数
    P->SmoothType = smooth_type;
    P->JacobiFactor = jacobi_factor;
    P->SORType = sor_type;
    P->SORFactor = sor_factor;
    P->ILUMaxFillins = ilu_maxfil;
    P->ILUDropTolerance = ilu_droptol;
    P->ChebyshevEigenRatio = chebyshev_eigen_ratio;
    P->PrintStats = print_stats;

    std::cout << "\n# AMG Setup Phase\n";
    std::cout << "--------------------------------------------------\n";

    P->Setup(*A);

    std::cout << "\n# BiCGStab Solve Phase\n";
    std::cout << "--------------------------------------------------\n";

    auto solverCG = std::make_unique<YHAMG::ParSolverBCGS>(max_iters, tolerance, print_stats);

    int iter;
    double relres;

    {
        PROFILE_SCOPE("Jacobi-BiCGStab Solve CPU");
        (*solverCG)(*A, *P, *b, *x, iter, relres);
    }

    MPI_Finalize();
    return 0;
}
